# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

import numpy as np
import torch

from fairseq.data import data_utils, FairseqDataset


def collate_3d_tokens(values, pad_idx, eos_idx=None, left_pad=False, move_eos_to_beginning=False):
    """Convert a list of padded 2d tensors into a padded 3d tensor."""
    number_entity = max(v.size(0) for v in values)
    size_entity = max(v.size(1) for v in values)
    res = values[0].new(len(values), number_entity, size_entity).fill_(pad_idx)

    def copy_tensor(src, dst):
        assert dst.numel() == src.numel()
        if move_eos_to_beginning:
            assert src[-1] == eos_idx
            dst[0] = eos_idx
            dst[1:] = src[:-1]
        else:
            dst.copy_(src)

    for i, v in enumerate(values):
        copy_tensor(v, res[i][number_entity - v.size(0):, size_entity - v.size(1):] if left_pad else res[i][:v.size(0), :v.size(1)])
    return res


def collate(
    samples, pad_idx, eos_idx, left_pad_source=True, left_pad_target=False,
    input_feeding=True, sep_idx=None, tgt_eos_idx=None, tgt_pad_idx=None, use_context=None
):
    if len(samples) == 0:
        return {}

    def merge(key, _pad_idx, _eos_idx, left_pad, move_eos_to_beginning=False):
        return data_utils.collate_tokens(
            [s[key] for s in samples],
            _pad_idx, _eos_idx, left_pad, move_eos_to_beginning,
        )

    def merge_3d_tensor(key, _pad_idx, left_pad, move_eos_to_beginning=False):
        return collate_3d_tokens(
            [s[key] for s in samples],
            _pad_idx, eos_idx, left_pad, move_eos_to_beginning,
        )

    id = torch.LongTensor([s['id'] for s in samples])
    if len(samples[0]['source'].size()) == 2:
        src_tokens = merge_3d_tensor('source', pad_idx, left_pad=left_pad_source)
        src_lengths = (src_tokens != pad_idx).sum(dim=2, keepdim=False)
    else:
        src_tokens = merge('source', pad_idx, eos_idx, left_pad=left_pad_source)
        # sort by descending source length
        src_lengths = torch.LongTensor([s['source'].numel() for s in samples])

    # entity
    entity_tokens = merge_3d_tensor('entity', pad_idx, left_pad=left_pad_source)
    entity_lengths = (entity_tokens != pad_idx).sum(dim=2, keepdim=False)

    prev_output_tokens = None
    context = None
    target = None
    tgt_event = None
    tgt_entity = None
    tgt_prompt_mask = None

    if samples[0].get('target', None) is not None:
        target = merge('target', tgt_pad_idx, tgt_eos_idx, left_pad=left_pad_target)
        target_lengths = [len(s['target']) for s in samples]
        ntokens = sum(target_lengths)

        if samples[0].get('tgt_event', None) is not None:
            # event label and entity label
            tgt_event = merge('tgt_event', -1, tgt_eos_idx, left_pad=left_pad_target)
            tgt_entity = merge('tgt_entity', -1, tgt_eos_idx, left_pad=left_pad_target)

            assert target.size() == tgt_event.size() == tgt_entity.size()

        # for pad == eos in bpe, so mask by length
        tgt_prompt_mask = torch.ones_like(target).type(torch.uint8)
        for i in range(len(samples)):
            tgt_prompt_mask[i][target_lengths[i]:] = 0

        # # prompt mask
        # if sep_idx is not None:
        #     has_mask = (target == sep_idx)
        #     tgt_prompt_len = has_mask.nonzero() + 1
        #     assert len(samples) == len(tgt_prompt_len)
        #     for i in range(len(samples)):
        #         tgt_prompt_mask[i][:tgt_prompt_len[i][1]] = 0
        #
        #     # update ntokens
        #     ntokens = ntokens - tgt_prompt_mask.eq(0).sum().item()
        if use_context and sep_idx is not None:
            def get_len(item):
                for i in range(len(item)):
                    if item[i: i+ len(sep_idx)].tolist() == sep_idx:
                        return i + len(sep_idx)
                return len(item)
            context_length = list(map(get_len, [s['target'] for s in samples]))
            max_len = max(context_length)
            context = target.new(target.size(0), max_len).long().fill_(tgt_pad_idx)
            for i in range(len(samples)):
                context[i, :context_length[i]] = target[i, :context_length[i]]

        if input_feeding:
            # we create a shifted version of targets for feeding the
            # previous output token(s) into the next decoder step
            prev_output_tokens = merge(
                'target', tgt_pad_idx, tgt_eos_idx,
                left_pad=left_pad_target,
                move_eos_to_beginning=True,
            )
    else:
        ntokens = sum((s['source'] != pad_idx).sum() for s in samples)

    batch = {
        'id': id,
        'nsentences': len(samples),
        'ntokens': ntokens,
        'net_input': {
            'src_tokens': src_tokens,
            'src_lengths': src_lengths,
            'entity_tokens': entity_tokens,
            'entity_lengths': entity_lengths,
        },
        'target': target,
        'tgt_event': tgt_event,
        'tgt_entity': tgt_entity,
        'tgt_prompt_mask': tgt_prompt_mask,
    }
    if prev_output_tokens is not None:
        batch['net_input']['prev_output_tokens'] = prev_output_tokens
    if context is not None:
        batch['context'] = context
    return batch


class HierEventStoryDataset(FairseqDataset):
    """
    A pair of torch.utils.data.Datasets.

    Args:
        src (torch.utils.data.Dataset): source dataset to wrap
        src_sizes (List[int]): source sentence lengths
        src_dict (~fairseq.data.Dictionary): source vocabulary
        tgt (torch.utils.data.Dataset, optional): target dataset to wrap
        tgt_sizes (List[int], optional): target sentence lengths
        tgt_dict (~fairseq.data.Dictionary, optional): target vocabulary
        left_pad_source (bool, optional): pad source tensors on the left side
            (default: True).
        left_pad_target (bool, optional): pad target tensors on the left side
            (default: False).
        max_source_positions (int, optional): max number of tokens in the
            source sentence (default: 1024).
        max_target_positions (int, optional): max number of tokens in the
            target sentence (default: 1024).
        shuffle (bool, optional): shuffle dataset elements before batching
            (default: True).
        input_feeding (bool, optional): create a shifted version of the targets
            to be passed into the model for input feeding/teacher forcing
            (default: True).
        remove_eos_from_source (bool, optional): if set, removes eos from end
            of source if it's present (default: False).
        append_eos_to_target (bool, optional): if set, appends eos to end of
            target if it's absent (default: False).
    """

    def __init__(
        self, src, src_sizes, src_dict,
        tgt=None, tgt_sizes=None, tgt_dict=None,
        entity_dataset=None, tgt_event_label_dataset=None, tgt_entity_label_dataset=None,
        left_pad_source=True, left_pad_target=False,
        max_source_positions=1024, max_target_positions=1024,
        shuffle=True, input_feeding=True, remove_eos_from_source=False, append_eos_to_target=False,
        use_context=False, truncate=False, flat_event=False
    ):
        # if tgt_dict is not None:
        #     assert src_dict.pad() == tgt_dict.pad()
        #     assert src_dict.eos() == tgt_dict.eos()
        #     assert src_dict.unk() == tgt_dict.unk()
        self.src = src
        self.tgt = tgt
        self.src_sizes = np.array(src_sizes)
        self.tgt_sizes = np.array(tgt_sizes) if tgt_sizes is not None else None
        self.src_dict = src_dict
        self.tgt_dict = tgt_dict
        self.left_pad_source = left_pad_source
        self.left_pad_target = left_pad_target
        self.max_source_positions = max_source_positions
        self.max_target_positions = max_target_positions
        self.shuffle = shuffle
        self.input_feeding = input_feeding
        self.remove_eos_from_source = remove_eos_from_source
        self.append_eos_to_target = append_eos_to_target
        self.entity = entity_dataset
        self.tgt_event = tgt_event_label_dataset
        self.tgt_entity = tgt_entity_label_dataset

        self.use_context = use_context
        self.truncate = truncate
        self.flat_event = flat_event

    def __getitem__(self, index):
        src_item = self.src[index]
        entity_item = self.entity[index]

        tgt_item = self.tgt[index] if self.tgt is not None else None
        tgt_event_item = self.tgt_event[index] if self.tgt_event is not None else None
        tgt_entity_item = self.tgt_entity[index] if self.tgt_entity is not None else None
        # Append EOS to end of tgt sentence if it does not have an EOS and remove
        # EOS from end of src sentence if it exists. This is useful when we use
        # use existing datasets for opposite directions i.e., when we want to
        # use tgt_dataset as src_dataset and vice versa
        if self.append_eos_to_target:
            eos = self.tgt_dict.eos() if self.tgt_dict else self.src_dict.eos()
            if self.tgt and self.tgt[index][-1] != eos:
                tgt_item = torch.cat([self.tgt[index], torch.LongTensor([eos])])

        # if self.remove_eos_from_source:
        #     eos = self.src_dict.eos()
        #     if self.src[index][-1] == eos:
        #         src_item = self.src[index][:-1]

        # for memory reason, only use 100 events
        event_num = 30
        if src_item.size(0) > event_num:
            src_item = src_item[:event_num]
            max_event_len = (src_item != self.src_dict.pad()).sum(dim=1).max().item()
            src_item = src_item[:, :max_event_len]

            if self.tgt_event is not None:
                event_mask = (tgt_event_item <= event_num)
                # select eos token
                event_mask[-1] = 1
                tgt_event_item = tgt_event_item[event_mask]
                tgt_event_item[-1] = tgt_event_item[-2]
                if self.tgt is not None:
                    tgt_item = tgt_item[event_mask]
                if self.tgt_entity is not None:
                    tgt_entity_item = tgt_entity_item[event_mask]
                assert tgt_event_item.size() == tgt_item.size() == tgt_entity_item.size()

        if self.flat_event:
            src_item = src_item[src_item.ne(self.src_dict.pad())]

        return {
            'id': index,
            'source': src_item,
            'target': tgt_item,
            'entity': entity_item,
            'tgt_event': tgt_event_item,
            'tgt_entity': tgt_entity_item,
        }

    def __len__(self):
        return len(self.src)

    def collater(self, samples):
        """Merge a list of samples to form a mini-batch.

        Args:
            samples (List[dict]): samples to collate

        Returns:
            dict: a mini-batch with the following keys:

                - `id` (LongTensor): example IDs in the original input order
                - `ntokens` (int): total number of tokens in the batch
                - `net_input` (dict): the input to the Model, containing keys:

                  - `src_tokens` (LongTensor): a padded 2D Tensor of tokens in
                    the source sentence of shape `(bsz, src_len)`. Padding will
                    appear on the left if *left_pad_source* is ``True``.
                  - `src_lengths` (LongTensor): 1D Tensor of the unpadded
                    lengths of each source sentence of shape `(bsz)`
                  - `prev_output_tokens` (LongTensor): a padded 2D Tensor of
                    tokens in the target sentence, shifted right by one position
                    for input feeding/teacher forcing, of shape `(bsz,
                    tgt_len)`. This key will not be present if *input_feeding*
                    is ``False``. Padding will appear on the left if
                    *left_pad_target* is ``True``.

                - `target` (LongTensor): a padded 2D Tensor of tokens in the
                  target sentence of shape `(bsz, tgt_len)`. Padding will appear
                  on the left if *left_pad_target* is ``True``.
        """
        if self.tgt_dict.index('<SEP>') != self.tgt_dict.unk():
            sep_idx = self.tgt_dict.index('<SEP>')
        elif getattr(self.tgt_dict, 'enc', False):
            sep_idx = self.tgt_dict.enc.encode(' <SEP>')
        else:
            sep_idx = None
        tgt_eos_idx = self.tgt_dict.eos()
        tgt_pad_idx = self.tgt_dict.pad()
        return collate(
            samples, pad_idx=self.src_dict.pad(), eos_idx=self.src_dict.eos(),
            left_pad_source=self.left_pad_source, left_pad_target=self.left_pad_target,
            input_feeding=self.input_feeding, sep_idx=sep_idx, tgt_eos_idx=tgt_eos_idx, tgt_pad_idx=tgt_pad_idx,
            use_context=self.use_context
        )

    def num_tokens(self, index):
        """Return the number of tokens in a sample. This value is used to
        enforce ``--max-tokens`` during batching."""
        return max(self.src_sizes[index], self.tgt_sizes[index] if self.tgt_sizes is not None else 0)

    def size(self, index):
        """Return an example's size as a float or tuple. This value is used when
        filtering a dataset with ``--max-positions``."""
        return (self.src_sizes[index], self.tgt_sizes[index] if self.tgt_sizes is not None else 0)

    def ordered_indices(self):
        """Return an ordered list of indices. Batches will be constructed based
        on this order."""
        if self.shuffle:
            indices = np.random.permutation(len(self))
        else:
            indices = np.arange(len(self))
        if self.tgt_sizes is not None:
            indices = indices[np.argsort(self.tgt_sizes[indices], kind='mergesort')]
        return indices[np.argsort(self.src_sizes[indices], kind='mergesort')]

    @property
    def supports_prefetch(self):
        return (
            getattr(self.src, 'supports_prefetch', False)
            and (getattr(self.tgt, 'supports_prefetch', False) or self.tgt is None)
            and getattr(self.entity, 'supports_prefetch', False)
            and (getattr(self.tgt_event, 'supports_prefetch', False) or self.tgt is None)
            and (getattr(self.tgt_entity, 'supports_prefetch', False) or self.tgt is None)
        )

    def prefetch(self, indices):
        self.src.prefetch(indices)
        self.entity.prefetch(indices)
        if self.tgt is not None:
            self.tgt.prefetch(indices)
            self.tgt_event.prefetch(indices)
            self.tgt_entity.prefetch(indices)

